Kth Largest Element in a Stream

问题描述（难度简单）

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Your KthLargest class will have a constructor which accepts an integer k and an integer array nums, which contains initial elements from the stream. For each call to the method KthLargest.add, return the element representing the kth largest element in the stream.

Example:

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

Note:
You may assume that nums‘ length ≥ k-1 and k ≥ 1.

方法一：排序

维护一个k大小的数组，排序时间复杂度klogk，每次进入新元素判断当前k个有序数组中最小值，小于最小值不进入，大于最小值进入。整体时间复杂度nklogk。

方法一代码

package FindKLargestNumber;

import java.util.ArrayList;
import java.util.Collections;

/**
 * @autor yeqiaozhu.
 * @date 2019-05-09
 */
public class KthLargestTwo {
    private final int k;
    private final ArrayList<Integer> array = new ArrayList<>();

    public KthLargestTwo(int k, int[] nums) {
        this.k = k;
        int numsLength = nums.length;

        for (int i = 0; i < numsLength; i++) {
            if (i > k-1) {
                Collections.sort(array);
                add(nums[i]);
            } else {
                array.add(nums[i]);
            }
        }
    }

    public int add(int val) {
        if (array.size() < k) {
            array.add(val);
        } else if (val > array.get(0)) {
            array.set(0, val);
        }
        Collections.sort(array);
        return array.get(0);
    }
}

方法二：优先队列

单纯排序的算法klogk的时间复杂度还有优化的空间，如果是最小堆的话搜素插入的时间只需要logk的复杂度，这里我们增加一个优先级队列的版本。优先级队列默认是一个堆的数据结构，堆这种数据结构是一个完全二叉树，满足每个子节点都要比父节点要小（小顶堆）或者大（大顶堆）。所以堆顶永远是k个数当中最小的，只需要比较堆顶的元素即可。同样的搜索插入的复杂度是logk。

方法二代码

package FindKLargestNumber;

import java.util.Arrays;
import java.util.PriorityQueue;

/**
 * Method 1:采用优先级队列（最小堆）的方式实现，nlogk时间复杂度，k的空间复杂度
 * Method 2:采用排序的方式，nklogk时间复杂度，k的空间复杂度
 * @autor yeqiaozhu.
 * @date 2019-05-09
 */
public class KthLargest {

    private final PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();
    private final int k;

    /**
     * 构建一个有k个元素的堆
     *
     * @param k
     * @param nums
     */
    public KthLargest(int k, int[] nums) {
        this.k = k;
        Arrays.stream(nums).forEach(integer ->
                add(integer));
    }

    public int add(int val) {
        //如果优先级队列中数量小于k，直接往最小堆中追加
        if (priorityQueue.size() < k) {
            priorityQueue.offer(val);
        } else if (priorityQueue.peek() < val) {//如果当前追加元素大于堆顶的元素，将堆顶的元素poll滚出，然后将当前的值offer进入堆
            priorityQueue.poll();
            priorityQueue.offer(val);
        }
        return priorityQueue.peek();
    }


    public static void main(String[] args) {
        int k = 3;
        int[] arr = {4, 5, 8, 2};
        KthLargest kthLargest = new KthLargest(3, arr);
        System.out.println(kthLargest.add(3));   // returns 4
        System.out.println(kthLargest.add(5));   // returns 5
        System.out.println(kthLargest.add(10));  // returns 5
        System.out.println(kthLargest.add(9));   // returns 8
        System.out.println(kthLargest.add(4));   // returns 8
    }
}

总结

本质上两种方式都是维护一个k大小的数据结构，优先级队列在搜索插入的时间复杂度上具备优势，只需要logk的时间复杂度，优先级队列本质上是一个小顶堆。我们可以利用优先级队列去优化方法一。